import logging
from typing import List
import requests
import aiohttp
import asyncio

from langchain.embeddings.base import Embeddings
from pydantic import BaseModel

logger = logging.getLogger(__name__)
import settings

if settings.ONLINE:
    # base_url = "http://172.30.0.22:9001"  #jd 已迁移到其他分支
    base_url = "http://172.24.0.8:9001"  #hs
else:
    # base_url = "http://116.196.97.15:9001"  #jd 已迁移到其他分支
    base_url = "http://101.126.35.198:9001"  #hs

class LocalEmbeddings(BaseModel, Embeddings):
    async def async_embedding(self, json):
        async with aiohttp.ClientSession() as session:
            async with session.post(f"{base_url}/embedding",
                                    headers={'Content-Type': 'application/json'},
                                    json=json) as response:
                res = await response.json()
                return res


    async def async_embed_documents(self, texts, chunk_size=16):
        text_in_chunks = [
            texts[i: i + chunk_size]
            for i in range(0, len(texts), chunk_size)
        ]
        batched_embeddings = []
        for chunk in text_in_chunks:
            resp = await self.async_embedding({"input": chunk})
            batched_embeddings.extend([i for i in resp["embedding"]])

        return batched_embeddings


    async def async_embed_query(self, text: str) -> List[float]:
        resp = await self.async_embedding({"input": text})
        return resp["embedding"]


    def _embedding(self, json: object) -> dict:
        resp = requests.post(
            f"{base_url}/embedding",
            headers={
                "Content-Type": "application/json",
            },
            json=json,
        )
        return resp.json()


    def embed_documents(self, texts: List[str], chunk_size=16) -> List[List[float]]:
        text_in_chunks = [
            texts[i : i + chunk_size]
            for i in range(0, len(texts), chunk_size)
        ]
        batched_embeddings = []
        for chunk in text_in_chunks:
            resp = self._embedding({"input": chunk})
            batched_embeddings.extend([i for i in resp["embedding"]])

        return batched_embeddings

    def embed_query(self, text: str) -> List[float]:
        resp = self._embedding({"input": text})

        return resp["embedding"]


if __name__ == "__main__":
    embd = LocalEmbeddings()
    text = ["我是中国人", "我是中国人"]
    print(asyncio.run(embd.async_embed_documents(text)))
